#pragma once
#include <iostream>
#include <assert.h>
#include <string.h>
#include <vector>

using std::cout;
using std::endl;

namespace real
{
	enum Color
	{
		RED,
		BLACK
	};

	template<class T>
	class RBTreeNode
	{
	public:
		Color _col;
		RBTreeNode<T>* _left;
		RBTreeNode<T>* _parent;
		RBTreeNode<T>* _right;
		T _data;
	// public:
		//friend class RBTree<K,V>;
		RBTreeNode(const T& data)
			: _col(RED)
			, _parent(nullptr)
			, _left(nullptr)
			, _right(nullptr)
			, _data(data)
		{}

		~RBTreeNode()
		{
			_parent = nullptr;
			_left = nullptr;
			_right = nullptr;
		}
	};

	template <class T, class Ref, class Ptr>
	struct RBTreeIterator
	{
		using Node = RBTreeNode<T> ;
		using Self = RBTreeIterator<T, Ref, Ptr> ;

		Node* _node;
		Node* _root;

		RBTreeIterator(Node* node, Node* root)
			:_node(node)
			,_root(root)
		{}

		Self operator++()
		{
			if (_node->_right)
			{
				Node* min = _node->_right;
				while (min->_left)
				{
					min = min->_left;
				}
				_node = min;
			}
			else 
			{
				Node* par = _node->_parent;
				Node* cur = _node;
				while (par && cur == par->_right)
				{
					cur = par;
					par = cur->_parent;
				}
				_node = par;
			}
			return *this;
		}

		Self operator--() {
			if(!_node)  // --end()
			{
				Node* max = _root;
				while (max && max->_right) {
					max = max->_right;
				}
				_node = max;
			}
			else if (_node->left) // 左子树不为空，中序左子树最后一个
			{
				Node* rightMost = _node->_left->right;
				while (rightMost && rightMost->_right) {
					rightMost = rightMost->_right;
				}
				_node = rightMost;
			}
			else  //
			{
				Node* parent = _node->_parent;
				while (parent && _node == parent->_left) {
					_node = parent;
					parent = parent->_parent;
				}
				_node = parent;
			}
			return *this;
		}

		Ref operator*()
		{
			return _node->_data;
		}

		Ptr operator->()
		{
			return &(_node->_data);
		}

		bool operator==(const Self& rhs) const
		{
			return _node == rhs._node;
		}

		bool operator!=(const Self& rhs) const {
			return _node != rhs._node;
		}
	};

	template<class K, class T, class KeyOfT>
	class RBTree
	{
	public:
		using iterator = RBTreeIterator<T, T&, T*>;
		using const_iterator = RBTreeIterator<T, const T&, const T*>;

		iterator begin()
		{
			Node* left = _root;
			while (left && left->_left) {
				left = left->_left;
			}
			return iterator(left, _root);
		}

		iterator end()
		{
			return iterator(nullptr, _root);
		}

		const_iterator begin() const
		{
			Node* left = _root;
			while (left && left->_left) {
				left = left->_left;
			}
			return const_iterator(left, _root);
		}

		const_iterator end() const
		{
			return const_iterator(nullptr, _root);
		}


	private:
		using Node = RBTreeNode<T>;
		Node* _root = nullptr;

		void _InOrder(Node* root)
		{
			if (!root) return;
			_InOrder(root->_left);
			cout << (root->_data) << "->";
			_InOrder(root->_right);
		}

		bool _IsBalancedTree(Node* root, int black_num, int ref_num)
		{
			if (root == nullptr)
			{
				if (black_num == ref_num)
				{
					return true;
				}
				return false;
			}

			if (root->_parent)
			{
				if (root->_parent->_col == RED && root->_col == RED)
				{
					cout << "检查到连续的红色节点" << endl;
					return false;
				}
			}

			if (root->_col == BLACK)
				++black_num;
			return _IsBalancedTree(root->_left, black_num, ref_num) && 
					_IsBalancedTree(root->_right, black_num, ref_num);
		}

		int _Height(Node* root)
		{
			if (!root) return 0;
			return 1 + std::max(_Height(root->_left), _Height(root->_right));
		}
	public:
		RBTree()
		{
			_root = nullptr;
		}

		iterator Find(const K& key) {
			Node* cur = _root;
			while (cur) {
				if (cur->_key == key)
					return iterator(cur, _root);
				if (cur->_key < key)
					cur = cur->_left;
				else
					cur = cur->_right;
			}
			return iterator(nullptr, _root);
		}

		std::pair<iterator, bool> Insert(const T& kv)
		{
			if (_root == nullptr)
			{
				_root = new Node(kv);
				_root->_col = BLACK;
				// return true;
				return std::make_pair(iterator(_root, _root), true);
			}
			KeyOfT kot;
			Node* cur = _root;
			Node* parent = nullptr;
			while (cur)
			{
				parent = cur;
				if (kot(cur->_data) < kot(kv))
				{
					cur = cur->_right;
				}
				else if (kot(cur->_data) >  kot(kv))
				{
					cur = cur->_left;
				}
				else
				{
					// return false;
					return std::make_pair(iterator(cur, _root), false);
				}
			}
			
			Node* newnode = new Node(kv);
			if (kot(parent->_data) < kot(kv))
			{
				parent->_right = newnode;
			}
			else
			{
				parent->_left = newnode;
			}
			newnode->_parent = parent;
			
			cur = newnode;

			while(cur->_parent)
			{
				parent = cur->_parent;
				if (parent->_col == BLACK) break;

				// parent->_col == RED  //Adjust RB Tree
				Node* grandpa = parent->_parent;
				if (parent == grandpa->_left)
				{
					//     g
					//   f   u
					// c
					Node* unc = grandpa->_right;
					if (unc && unc->_col == RED)
					{
						unc->_col = BLACK;
						parent->_col = BLACK;
						grandpa->_col = RED;
						cur = grandpa;
						continue;
					}
					else // unc not valid || unc->_col == BLACK 必定不是新增的,由上次循环产生的情况
					{
						//     g           g
						//   f   u       f   
						// c	       c
						if (cur == parent->_left)
						{
							
							RotateR(grandpa);
							cur->_col = RED;
							parent->_col = BLACK;
							grandpa->_col = RED;
							break;
						}
						else
						{
							//     g            g            c
							//   f   u  -->   c   u  -->  f     g
							//	   c		f					   u
							// 左右双旋
							// 既然cur是上次循环变上来的, 那么cur的child必定是BLACK, parent是红, parent的左子树也是BLACK
							// 所以cur作为新的root当BLACK, parent和grandpa当RED,每条分支黑色的数量不变
							RotateLR(grandpa);
							cur->_col = BLACK;
							parent->_col = RED;
							grandpa->_col = RED;
							break;
						}
					}
				}
				else  // parent == grandpa->_right
				{
					//     g
					//   u   f
					//		   c
					Node* unc = grandpa->_left;
					if (unc && unc->_col == RED)
					{
						parent->_col = BLACK;
						unc->_col = BLACK;
						grandpa->_col = RED;
						cur = grandpa;
						continue;
					}
					else  // unc not valid || unc->_col == BLACK 必定不是新增的,由上次循环产生的情况
					{
						//     g           g
						//   u   f           f   
						// 	       c          c
						if (cur == parent->_right)
						{

							RotateL(grandpa);
							cur->_col = RED;
							parent->_col = BLACK;
							grandpa->_col = RED;
							break;
						}
						else
						{
							//     g
							//   u   f
							//	    c
							// 右左双旋
							// 既然cur是上次循环变上来的, 那么cur的child必定是BLACK, parent是红, parent的左子树也是BLACK
							// 所以cur作为新的root当BLACK, parent和grandpa当RED,每条分支黑色的数量不变
							RotateRL(grandpa);
							cur->_col = BLACK;
							parent->_col = RED;
							grandpa->_col = RED;
							break;
						}
					}
				}
			}
			_root->_col = BLACK;
			// return true;
			return std::make_pair(iterator(newnode, _root), true);
		}

		void RotateL(Node* root)
		{
			Node* pparent = root->_parent;
			Node* subR = root->_right;
			Node* subRL = subR->_left;

			root->_right = subRL;
			if (subRL)
				subRL->_parent = root;
			root->_parent = subR;
			subR->_left = root;
			if (pparent)
			{
				if (root == pparent->_left)
					pparent->_left = subR;
				else
					pparent->_right = subR;
				subR->_parent = pparent;
			}
			else
			{
				subR->_parent = nullptr;
				_root = subR;
			}
		}

		void RotateR(Node* root)
		{
			Node* pparent = root->_parent;
			Node* subL = root->_left;
			Node* subLR = subL->_right;

			root->_left = subLR;
			if (subLR)
				subLR->_parent = root;
			root->_parent = subL;
			subL->_right = root;
			if (pparent)
			{
				if (root == pparent->_left)
					pparent->_left = subL;
				else
					pparent->_right = subL;
				subL->_parent = pparent;
			}
			else
			{
				subL->_parent = nullptr;
				_root = subL;
			}
		}

		void RotateLR(Node* root)
		{
			RotateL(root->_left);
			RotateR(root);

		}

		void RotateRL(Node* root)
		{
			RotateR(root->_right);
			RotateL(root);
		}

		void InOrder()
		{
			_InOrder(_root);
			cout << endl;
		}

		bool IsBalancedTree()
		{
			int ref_num = 0;
			Node* cur = _root;
			while (cur)
			{
				if (cur->_col == BLACK)
					++ref_num;
				cur = cur->_left;
			}
			return _IsBalancedTree(_root, 0, ref_num);
		}

		int Height()
		{
			return _Height(_root);
		}
	};
}
